import matplotlib.pyplot as plt

import pickle
import pandas as pd
import numpy as np
from datetime import datetime, timedelta

if __name__ == '__main__':
    option = int(input('可视化选择：'))
    # 按照国内国外分类
    if option == 1:
        vid = pickle.load(file=open('../PickledFile/video.pickled', mode='rb'))
        dms = pickle.load(file=open('../PickledFile/dm_predicted(1).pickled', mode='rb'))  # type: pd.DataFrame
        cts = pickle.load(file=open('../PickledFile/comments_predicted (1).pickled', mode='rb'))  # type: pd.DataFrame
        dms = dms.reset_index()
        time_tmp = []

        for i in dms.values:
            av = i[0]
            v = vid.loc[av]
            time_tmp.append(v['upload_time'])
        dms['comment_time'] = time_tmp
        dms = dms.set_index(keys=['video_id'])

        begin = datetime(year=2020, month=1, day=15)
        end = datetime(year=2020, month=5, day=31)

        x = []
        y1 = []
        y2 = []

        while begin < end:
            tmp1 = cts[(begin <= cts['comment_time']) & (cts['comment_time'] <= (begin + timedelta(days=5)))]
            tmp2 = dms[(begin <= dms['comment_time']) & (dms['comment_time'] <= (begin + timedelta(days=5)))]
            a = int(tmp1[tmp1['res'] == 0].size) + int(tmp2[tmp2['res'] == 0].size)
            b = int(tmp1[tmp1['res'] == 1].size) + int(tmp2[tmp2['res'] == 1].size)
            c = int(tmp1[tmp1['res'] == 2].size) + int(tmp2[tmp2['res'] == 2].size)
            print(a, b, c)
            x.append(begin)
            y1.append(b / (a + b + c))
            y2.append(c / (a + b + c))
            begin = begin + timedelta(days=5)

        total = dms.size + cts.size
        domestic = dms[dms['res'] == 1].size + cts[cts['res'] == 1].size
        abroad = dms[dms['res'] == 2].size + cts[cts['res'] == 2].size

        abroad = abroad / total
        domestic = domestic / total
        other = 1 - abroad - domestic

    # 按照内容性质：调侃、鼓励、严肃进行分类
    if option == 2:
        vid = pickle.load(file=open('../PickledFile/video.pickled', mode='rb'))
        dms = pickle.load(file=open(file='../PickledFile/dm_predicted(3).pickled', mode='rb'))  # type: pd.DataFrame
        cts = pickle.load(file=open(file='../PickledFile/comments_predicted (3).pickled', mode='rb'))  # type: pd.DataFrame
        dms = dms.reset_index()
        time_tmp = []

        for i in dms.values:
            av = i[0]
            v = vid.loc[av]
            time_tmp.append(v['upload_time'])
        dms['comment_time'] = time_tmp
        dms = dms.set_index(keys=['video_id'])

        begin = datetime(year=2020, month=1, day=15)
        end = datetime(year=2020, month=5, day=31)

        x, y0, y1, y2, y3 = [], [], [], [], []
        while begin < end:
            tmp1 = dms[(dms['comment_time'] >= begin) & (dms['comment_time'] < (begin + timedelta(days=5)))]
            tmp2 = cts[(cts['comment_time'] >= begin) & (cts['comment_time'] < (begin + timedelta(days=5)))]
            a = tmp1[tmp1['res'] == 0].size + tmp2[tmp2['res'] == 0].size
            b = tmp1[tmp1['res'] == 1].size + tmp2[tmp2['res'] == 1].size
            c = tmp1[tmp1['res'] == 2].size + tmp2[tmp2['res'] == 2].size
            d = tmp1[tmp1['res'] == 3].size + tmp2[tmp2['res'] == 3].size
            x.append(begin)
            y0.append(a / (a + b + c + d))
            y1.append(b / (a + b + c + d))
            y2.append(c / (a + b + c + d))
            y3.append(d / (a + b + c + d))
            begin = begin + timedelta(days=5)

        total = dms.size + cts.size
        ridicule = (dms[dms['res'] == 1].size + cts[cts['res'] == 1].size) / total
        encourage = (dms[dms['res'] == 2].size + cts[cts['res'] == 2].size) / total
        serious = (dms[dms['res'] == 3].size + cts[cts['res'] == 3].size) / total
        other = (total - ridicule - encourage - serious) / total

    # 杂七杂八的可视化
    if option == 3:
        vid = pickle.load(file=open(file='../PickledFile/video.pickled', mode='rb'))  # type: pd.DataFrame
        dms = pickle.load(file=open(file='../PickledFile/dm.pickled', mode='rb'))  # type: pd.DataFrame
        cts = pickle.load(file=open(file='../PickledFile/comments.pickled', mode='rb'))  # type: pd.DataFrame
        time_tmp = []
        dms = dms.reset_index()
        for i in dms.values:
            av = i[0]
            v = vid.loc[av]
            time_tmp.append(v['upload_time'])
        dms['comment_time'] = time_tmp
        dms = dms.set_index(keys=['video_id'])

        begin = datetime(year=2020, month=1, day=15)
        end = datetime(year=2020, month=5, day=31)

        x = []
        v_cnt = []
        cts_cnt = []
        dms_cnt = []

        while begin <= end:
            x.append(begin)
            v_cnt.append(vid[(vid['upload_time'] >= begin) &
                             (vid['upload_time'] < (begin + timedelta(days=5)))].size)
            cts_cnt.append(cts[(cts['comment_time'] >= begin) &
                               (cts['comment_time'] < (begin + timedelta(days=5)))].size)
            dms_cnt.append(dms[(dms['comment_time'] >= begin) &
                               (dms['comment_time'] < (begin + timedelta(days=5)))].size)
            begin = begin + timedelta(days=5)
